'''
Description: 
version: 
Author: 
Date: 2023-12-19 11:14:06
LastEditors: 
LastEditTime: 2023-12-19 11:14:06
'''

import torch 

text = "This was a masterpiece. Not completely faithful to the books, but enthralling from beginning to end. Might be my favorite of the three."


from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

model_path = "output/text-classification/checkpoint-3"
tokenizer = AutoTokenizer.from_pretrained(model_path)
inputs = tokenizer(text, return_tensors="pt")


model = AutoModelForSequenceClassification.from_pretrained(model_path)
with torch.no_grad():
    logits = model(**inputs).logits


predicted_class_id = logits.argmax().item()

print(model.config.id2label[predicted_class_id])